"""
生产环境的Vanna AI数据查询服务 - 道麒项目
优化版本，包含配置管理、错误处理、日志记录等功能
"""

import os
import logging
from logging.handlers import RotatingFileHandler, TimedRotatingFileHandler
from typing import Optional, Dict, Any
from dataclasses import dataclass
from functools import wraps
import asyncio
from datetime import datetime
import pathlib
import nest_asyncio
import uvicorn
from typing import List

# 应用nest_asyncio以允许嵌套事件循环
nest_asyncio.apply()

# 导入Vanna相关模块
from vanna.integrations.ollama import OllamaLlmService
from vanna.tools import RunSqlTool
from vanna.integrations.mysql import MySQLRunner
from vanna.core.tool import ToolContext, ToolResult
from vanna.tools.agent_memory import SaveQuestionToolArgsTool, SearchSavedCorrectToolUsesTool, SaveTextMemoryTool
from vanna.integrations.chromadb import ChromaAgentMemory
from vanna.core.user import UserResolver, User, RequestContext
from vanna import Agent
from vanna.core.registry import ToolRegistry
from vanna.tools import VisualizeDataTool
from vanna.core.agent import AgentConfig
from vanna.servers.fastapi import VannaFastAPIServer
from vanna.core.enhancer import LlmContextEnhancer
from vanna.core.llm import LlmMessage
from vanna.capabilities.sql_runner import RunSqlToolArgs

# 配置日志
def setup_logging():
    """设置日志配置，包括控制台和文件输出"""
    # 创建logs目录（如果不存在）
    logs_dir = pathlib.Path("logs")
    logs_dir.mkdir(exist_ok=True)
    
    # 获取当前日期作为日志文件名的一部分
    current_date = datetime.now().strftime("%Y%m%d")
    log_file = logs_dir / f"vanna_prod_{current_date}.log"
    
    # 配置日志格式
    formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 创建根日志记录器
    root_logger = logging.getLogger()
    root_logger.setLevel(logging.INFO)
    
    # 清除现有的处理器
    for handler in root_logger.handlers[:]:
        root_logger.removeHandler(handler)
    
    # 添加控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(formatter)
    root_logger.addHandler(console_handler)
    
    # 添加文件处理器（按大小轮转）
    file_handler = RotatingFileHandler(
        log_file,
        maxBytes=10*1024*1024,  # 10MB
        backupCount=5,
        encoding='utf-8'
    )
    file_handler.setLevel(logging.INFO)
    file_handler.setFormatter(formatter)
    root_logger.addHandler(file_handler)
    
    # 添加按时间轮转的处理器（每天午夜轮转）
    time_handler = TimedRotatingFileHandler(
        logs_dir / "vanna_prod_daily.log",
        when='midnight',
        interval=1,
        backupCount=30,
        encoding='utf-8'
    )
    time_handler.setLevel(logging.INFO)
    time_handler.setFormatter(formatter)
    root_logger.addHandler(time_handler)
    
    # 记录日志系统启动信息
    logger = logging.getLogger(__name__)
    logger.info(f"日志系统已初始化，日志文件: {log_file}")

# 初始化日志系统
setup_logging()
logger = logging.getLogger(__name__)

# 常用表配置 - 道麒项目核心业务表
COMMON_TABLES = ['dr_crm_account', 'dr_crm_potential','dr_crm_sales_order','dr_crm_sales_order_product','dr_crm_quotes','dr_crm_sale_contract']

# 业务上下文：用于存储动态生成的建表语句
business_contexts: List[str] = []

def handle_errors(func):
    """错误处理装饰器，用于捕获和记录异常"""
    @wraps(func)
    async def wrapper(*args, **kwargs):
        try:
            return await func(*args, **kwargs)
        except Exception as e:
            logger.error(f"在执行 {func.__name__} 时发生错误: {str(e)}", exc_info=True)
            raise
    return wrapper


@dataclass
class DatabaseConfig:
    """数据库配置类"""
    host: str
    database: str
    user: str
    password: str
    port: int
    
    @classmethod
    def from_env(cls) -> 'DatabaseConfig':
        """从环境变量创建数据库配置"""
        return cls(
            host=os.getenv('DB_HOST', '192.168.1.238'),
            database=os.getenv('DB_NAME', 'hdp_manage_dq'),
            user=os.getenv('DB_USER', 'root'),
            password=os.getenv('DB_PASSWORD', 'Hc$msql%0226.'),
            port=int(os.getenv('DB_PORT', '3306'))
        )


@dataclass
class LLMConfig:
    """LLM配置类"""
    model: str
    host: str
    
    @classmethod
    def from_env(cls) -> 'LLMConfig':
        """从环境变量创建LLM配置"""
        return cls(
            model=os.getenv('LLM_MODEL', 'gpt-oss:120b'),
            host=os.getenv('LLM_HOST', 'http://zh.daoreach.com:11434')
        )


@dataclass
class ServerConfig:
    """服务器配置类"""
    port: int
    max_tool_iterations: int
    temperature: float
    
    @classmethod
    def from_env(cls) -> 'ServerConfig':
        """从环境变量创建服务器配置"""
        return cls(
            port=int(os.getenv('SERVER_PORT', '8001')),
            max_tool_iterations=int(os.getenv('MAX_TOOL_ITERATIONS', '30')),
            temperature=float(os.getenv('LLM_TEMPERATURE', '0.7'))
        )


class ProductionEditableRunSqlTool(RunSqlTool):
    """生产环境SQL执行工具，增强错误处理和日志记录"""
    
    @handle_errors
    async def execute(self, context: ToolContext, args) -> ToolResult:
        """
        执行SQL查询，记录生成的SQL语句
        
        Args:
            context: 工具执行上下文
            args: SQL执行参数
            
        Returns:
            ToolResult: SQL执行结果
        """
        original_sql = args.sql
        logger.info(f"用户 {context.user.email} 执行SQL查询:\n{original_sql}")
        
        # 记录执行开始时间，用于性能监控
        import time
        start_time = time.time()
        
        # 执行SQL查询
        result = await super().execute(context, args)
        
        # 记录执行时间，用于性能分析
        execution_time = time.time() - start_time
        logger.info(f"SQL查询执行完成，耗时: {execution_time:.2f}秒")
        
        return result


class ProductionUserResolver(UserResolver):
    """生产环境用户解析器，增强安全性"""
    
    async def resolve_user(self, request_context: RequestContext) -> User:
        """
        解析用户信息，支持多种认证方式
        
        Args:
            request_context: 请求上下文
            
        Returns:
            User: 解析后的用户对象
        """
        # 尝试从Cookie获取用户邮箱（前端会话存储）
        user_email = request_context.get_cookie('vanna_email')
        
        # 如果Cookie中没有，尝试从Header获取（API调用）
        if not user_email:
            user_email = request_context.get_header('X-User-Email')
        
        # 如果都没有，使用默认访客用户（安全回退）
        if not user_email:
            user_email = 'guest@example.com'
            logger.warning("未找到用户信息，使用默认访客账户")
        
        # 根据邮箱确定用户组（简单权限控制）
        group = 'admin' if user_email == 'admin@example.com' else 'user'
        
        logger.info(f"用户 {user_email} 已认证，组: {group}")
        return User(id=user_email, email=user_email, group_memberships=[group])


class ProductionLlmContextEnhancer(LlmContextEnhancer):
    """
    生产环境LLM上下文增强器，优化表结构获取和缓存，增强业务表信息
    """
    
    def __init__(self, mysql_runner: MySQLRunner, agent_memory=None):
        """
        初始化LLM上下文增强器
        
        Args:
            mysql_runner: MySQL运行器实例，用于执行数据库查询
            agent_memory: 代理内存实例，用于存储和检索上下文信息
        """
        self.mysql_runner = mysql_runner
        self.agent_memory = agent_memory
        self.table_schema_cache = {}  # 表结构缓存，避免重复查询
        self.table_info_cache = {}    # 表信息缓存（包含注释、业务描述等）
        self.all_tables = []          # 数据库中所有可用表的列表
        self._initialized = False     # 初始化状态标志
        self._init_lock = asyncio.Lock()  # 异步锁，防止并发初始化
        
        # 预定义的业务表描述信息
        self.table_business_descriptions = {
            'dr_crm_account': '客户账户表，存储客户的基本信息、联系方式和账户状态',
            'dr_crm_potential': '销售机会表，记录潜在客户和销售线索的详细信息和跟进状态',
            'dr_crm_sales_order': '销售订单表，包含订单的基本信息、客户信息、订单状态和金额',
            'dr_crm_sales_order_product': '销售订单产品明细表，记录订单中的具体产品信息、数量和单价',
            'dr_crm_quotes': '报价单表，存储给客户的报价信息、报价状态和有效期',
            'dr_crm_sale_contract': '销售合同表，记录与客户签订的正式合同信息、条款和履行状态'
        }
    
    @handle_errors
    async def _ensure_initialized(self):
        """确保已初始化表结构信息，使用锁防止并发初始化
        
        采用双重检查锁定模式，确保在并发环境下只初始化一次
        """
        if not self._initialized:
            async with self._init_lock:
                if not self._initialized:
                    await self._get_all_tables()
                    await self._get_common_table_structures()
                    self._initialized = True
                    logger.info(f"表结构初始化完成，共加载 {len(self.table_schema_cache)} 个表结构")
    
    @handle_errors
    async def _get_all_tables(self):
        """获取数据库中的所有表名
        
        只查询配置在COMMON_TABLES中的表，避免加载不必要的表结构
        """
        if self.all_tables:
            return self.all_tables
            
        try:
            # 创建系统用户上下文，用于执行管理操作
            context = ToolContext(
                user=User(id="system", email="system@vanna.ai", group_memberships=['admin']),
                conversation_id="schema_loading",
                request_id="schema_request",
                agent_memory=self.agent_memory
            )
            
            # 构建查询语句，只获取配置的常用表
            tables_query = "SELECT table_name FROM information_schema.tables WHERE table_schema = '" + self.mysql_runner.database + "' AND table_name IN (" + ", ".join([f"'" + t + "'" for t in COMMON_TABLES]) + ")"
            result = await self.mysql_runner.run_sql(
                RunSqlToolArgs(sql=tables_query),
                context
            )
            
            if not result.empty:
                self.all_tables = result.iloc[:, 0].tolist()
                logger.info(f"获取到数据库表列表: {self.all_tables}")
            
        except Exception as e:
            logger.error(f"获取表列表时出错: {e}")
            raise
            
        return self.all_tables
    
    @handle_errors
    async def _get_common_table_structures(self) -> None:
        """获取常用表的结构
        
        预加载所有配置的常用表结构，提高后续查询响应速度
        """
        # 确保COMMON_TABLES中指定的表结构都已加载
        for table_name in COMMON_TABLES:
            # 如果表结构尚未缓存，则获取它
            if table_name not in self.table_schema_cache:
                await self._get_table_structure(table_name)
    
    @handle_errors
    async def _get_table_structure(self, table_name: str) -> str:
        """
        从数据库中获取表结构，包括表结构、注释和字段信息
        
        Args:
            table_name: 表名
            
        Returns:
            str: 表结构信息，包含完整的表定义和元数据
        """
        # 先检查缓存，避免重复查询
        if table_name in self.table_schema_cache:
            return self.table_schema_cache[table_name]
        
        try:
            # 创建系统用户上下文
            context = ToolContext(
                user=User(id="system", email="system@vanna.ai", group_memberships=['admin']),
                conversation_id="schema_loading",
                request_id="schema_request",
                agent_memory=self.agent_memory
            )
            
            # 获取表的完整信息（包括结构、注释、字段等）
            table_info = await self._get_table_info(table_name, context)
            
            if table_info:
                # 构建增强的表结构信息
                enhanced_schema = self._build_enhanced_table_schema(table_info)
                
                # 缓存结果，提高后续查询效率
                self.table_schema_cache[table_name] = enhanced_schema
                self.table_info_cache[table_name] = table_info
                
                logger.debug(f"获取表 {table_name} 增强结构成功")
                return enhanced_schema
            
        except Exception as e:
            logger.error(f"获取表 {table_name} 结构时出错: {e}")
        
        return ""
    
    async def _get_table_info(self, table_name: str, context: ToolContext) -> dict:
        """
        获取表的完整信息，包括表结构、注释、字段信息等
        
        Args:
            table_name: 表名
            context: 工具上下文
            
        Returns:
            dict: 包含表完整信息的字典
        """
        table_info = {
            'name': table_name,
            'database': self.mysql_runner.database,
            'create_sql': '',
            'table_comment': '',
            'business_description': self.table_business_descriptions.get(table_name, ''),
            'columns': []
        }
        
        try:
            # 1. 获取CREATE TABLE语句
            create_query = f"SHOW CREATE TABLE {table_name}"
            create_result = await self.mysql_runner.run_sql(
                RunSqlToolArgs(sql=create_query),
                context
            )
            if not create_result.empty and 'Create Table' in create_result.columns:
                table_info['create_sql'] = create_result['Create Table'].iloc[0]
            
            # 2. 获取表注释
            comment_query = f"""
            SELECT table_comment 
            FROM information_schema.tables 
            WHERE table_schema = '{self.mysql_runner.database}' 
            AND table_name = '{table_name}'
            """
            comment_result = await self.mysql_runner.run_sql(
                RunSqlToolArgs(sql=comment_query),
                context
            )
            if not comment_result.empty:
                table_info['table_comment'] = comment_result.iloc[0, 0] or ''
            
            # 3. 获取字段信息
            # columns_query = f"""
            # SELECT column_name, column_type, column_comment 
            # FROM information_schema.columns 
            # WHERE table_schema = '{self.mysql_runner.database}' 
            # AND table_name = '{table_name}' 
            # ORDER BY ordinal_position
            # """
            # columns_result = await self.mysql_runner.run_sql(
            #     RunSqlToolArgs(sql=columns_query),
            #     context
            # )
            # if not columns_result.empty:
            #     for _, row in columns_result.iterrows():
            #         table_info['columns'].append({
            #             'name': row['COLUMN_NAME'],
            #             'type': row['COLUMN_TYPE'],
            #             'comment': row['COLUMN_COMMENT'] or ''
            #         })
            
        except Exception as e:
            logger.error(f"获取表 {table_name} 详细信息时出错: {e}")
        
        return table_info
    
    def _build_enhanced_table_schema(self, table_info: dict) -> str:
        """
        构建增强的表结构描述，包含业务信息和字段说明
        
        Args:
            table_info: 表信息字典
            
        Returns:
            str: 格式化的增强表结构描述
        """
        table_name = table_info['name']
        database = table_info['database']
        
        # 构建增强的表描述
        enhanced_schema = []
        enhanced_schema.append(f"# 表信息: {database}.{table_name}")
        enhanced_schema.append("-" * 60)
        
        # 添加业务描述
        if table_info['business_description']:
            enhanced_schema.append(f"\n## 业务描述")
            enhanced_schema.append(table_info['business_description'])
        
        # 添加表注释
        if table_info['table_comment']:
            enhanced_schema.append(f"\n## 表注释")
            enhanced_schema.append(table_info['table_comment'])
        
        # # 添加字段信息表格
        # if table_info['columns']:
        #     enhanced_schema.append(f"\n## 字段信息")
        #     enhanced_schema.append(f"| {'字段名':<20} | {'数据类型':<20} | {'字段注释':<40} |")
        #     enhanced_schema.append(f"|{'-'*22}|{'-'*22}|{'-'*42}|")
            
        #     for column in table_info['columns']:
        #         enhanced_schema.append(f"| {column['name']:<20} | {column['type']:<20} | {column['comment'][:37]:<40} |")
        
        # # 添加建表语句
        # enhanced_schema.append(f"\n## 建表语句")
        # enhanced_schema.append(table_info['create_sql'])
        
        # 添加表使用说明
        enhanced_schema.append(f"\n## 使用说明")
        enhanced_schema.append(f"- 表名格式: {database}.{table_name}")
        enhanced_schema.append(f"- 请严格使用表中定义的字段名")
        enhanced_schema.append(f"- 根据业务需求可能需要与其他表进行JOIN操作")
        
        return "\n".join(enhanced_schema)

    @handle_errors
    async def generate_business_contexts(self, table_names: Optional[list[str]] = None) -> list[str]:
        """
        动态生成业务表的建表语句上下文，并返回列表

        说明：
        - 优先使用传入的 `table_names`；未指定时使用 COMMON_TABLES；若为空则回退到全部表
        - 每条上下文包含完整的 `CREATE TABLE` 语句（字段、类型、约束等）
        - 用于初始化代理内存，提供结构化的业务表信息
        """
        await self._ensure_initialized()

        names: list[str] = []
        if table_names:
            names = table_names
        elif COMMON_TABLES:
            names = COMMON_TABLES
        else:
            names = self.all_tables or []

        contexts: list[str] = []
        for name in names:
            schema = await self._get_table_structure(name)
            if schema:
                contexts.append(schema)

        # 若未能获取到指定表的结构，回退尝试数据库中的其他表（最多3个）
        if not contexts and self.all_tables:
            for name in self.all_tables:
                schema = await self._get_table_structure(name)
                if schema:
                    contexts.append(schema)
                if len(contexts) >= 3:
                    break

        return contexts
    

    @handle_errors
    async def enhance_system_prompt(
        self,
        system_prompt: str,
        user_message: str,
        user: User
    ) -> str:
        """
        增强系统提示词，添加表结构信息、业务描述和查询指导
        
        Args:
            system_prompt: 原始系统提示词
            user_message: 用户消息
            user: 用户对象
            
        Returns:
            str: 增强后的系统提示词
        """
        # 确保表信息已初始化
        await self._ensure_initialized()
        
        database_name = self.mysql_runner.database
        
        # 添加数据库基本信息
        enhanced_prompt = [system_prompt]
        enhanced_prompt.append("\n\n# 数据库和业务表信息")
        enhanced_prompt.append("=" * 60)
        enhanced_prompt.append(f"## 数据库信息")
        enhanced_prompt.append(f"- 数据库名称: {database_name}")
        enhanced_prompt.append(f"- 适用表名格式: {database_name}.表名")
        
        # 添加增强的业务表信息
        enhanced_prompt.append("\n## 核心业务表详细信息")
        enhanced_prompt.append("-" * 60)
        
        # 获取并添加所有常用表的增强结构信息
        table_count = 0
        for table_name in COMMON_TABLES:
            # 确保表结构已加载
            if table_name not in self.table_schema_cache:
                await self._get_table_structure(table_name)
            
            if table_name in self.table_schema_cache:
                enhanced_prompt.append("\n")
                enhanced_prompt.append(self.table_schema_cache[table_name])
                enhanced_prompt.append("\n" + "-" * 60)
                table_count += 1
        
        # 添加通用查询指导
        enhanced_prompt.append("\n## SQL查询指导")
        enhanced_prompt.append("- 请严格使用上述表结构中定义的表名和字段名")
        enhanced_prompt.append("- 根据业务需求选择合适的表，可能需要查询多个表并使用JOIN操作")
        enhanced_prompt.append("- 生成的SQL必须符合MySQL语法规范")
        enhanced_prompt.append("- 查询结果应包含有意义的列名别名，提高可读性")
        
        # 添加对数量查询的明确指导（针对中文用户的常见查询模式）
        if any(keyword in user_message.lower() for keyword in ['多少', '数量', '总数', 'count']):
            enhanced_prompt.append("\n## 数量查询特殊指导")
            enhanced_prompt.append("- 当用户询问数量相关问题时，请生成简单的COUNT查询")
            enhanced_prompt.append(f"- 正确示例: SELECT COUNT(*) AS total_count FROM {database_name}.表名")
            enhanced_prompt.append("- 不要添加额外的WHERE条件，除非用户明确指定了过滤条件")
            enhanced_prompt.append("- 为结果字段提供有意义的别名，如total_count、customer_count等")
        
        # 添加用户身份信息（用于权限控制）
        enhanced_prompt.append(f"\n## 用户信息")
        enhanced_prompt.append(f"- 当前用户: {user.email}")
        enhanced_prompt.append(f"- 用户组: {', '.join(user.group_memberships)}")
        
        enhanced_prompt.append("\n请根据用户的问题和提供的表结构信息，生成准确、高效、符合业务逻辑的SQL查询。")
        
        # 日志记录增强情况
        logger.info(f"已增强系统提示词，添加了 {table_count} 个表的详细结构信息")
        # 记录用户消息
        logger.info(f"用户消息: {user_message}")
        # 记录增强后的系统提示词
        logger.info(f"增强后的系统提示词: {enhanced_prompt}")
        
        return "\n".join(enhanced_prompt)

    @handle_errors
    async def enhance_user_messages(
        self,
        messages: list[LlmMessage],
        user: User
    ) -> list[LlmMessage]:
        """
        增强用户消息，添加用户上下文信息
        
        Args:
            messages: 原始消息列表
            user: 用户对象
            
        Returns:
            list[LlmMessage]: 增强后的消息列表，包含用户身份信息
        """
        messages.append(LlmMessage(
            role="system", 
            content=f"[系统上下文信息：当前登录用户为 {user.email}，所属组为 {', '.join(user.group_memberships)}。此信息仅用于身份验证和权限控制，不是用户查询条件。]"
        ))
        return messages


async def create_production_agent() -> Agent:
    """
    创建生产环境代理实例
    
    Returns:
        Agent: 配置好的代理实例
    """
    # 加载配置
    db_config = DatabaseConfig.from_env()
    llm_config = LLMConfig.from_env()
    server_config = ServerConfig.from_env()
    
    logger.info("正在初始化生产环境代理...")
    
    # 初始化LLM服务
    llm = OllamaLlmService(
        model=llm_config.model,
        host=llm_config.host
    )
    logger.info(f"LLM服务已初始化: {llm_config.model} @ {llm_config.host}")
    
    # 初始化数据库工具
    db_tool = ProductionEditableRunSqlTool(
        sql_runner=MySQLRunner(
            host=db_config.host,
            database=db_config.database,
            user=db_config.user,
            password=db_config.password,
            port=db_config.port
        )
    )
    logger.info(f"数据库连接已初始化: {db_config.user}@{db_config.host}:{db_config.port}/{db_config.database}")
    
    # 初始化代理内存
    agent_memory = ChromaAgentMemory(
        collection_name="vanna_memory",
        persist_directory="./chroma_db"
    )
    logger.info("代理内存已初始化")
    
    # 初始化用户解析器
    user_resolver = ProductionUserResolver()
    logger.info("用户解析器已初始化")
    
    # 注册工具
    tools = ToolRegistry()
    tools.register_local_tool(db_tool, access_groups=['admin', 'user'])
    tools.register_local_tool(SaveQuestionToolArgsTool(), access_groups=['admin'])
    tools.register_local_tool(SearchSavedCorrectToolUsesTool(), access_groups=['admin', 'user'])
    tools.register_local_tool(SaveTextMemoryTool(), access_groups=['admin', 'user'])
    tools.register_local_tool(VisualizeDataTool(), access_groups=['admin', 'user'])
    logger.info("工具注册完成")
    
    # 配置代理参数
    config = AgentConfig(
        max_tool_iterations=server_config.max_tool_iterations,
        stream_mode=True,
        temperature=server_config.temperature,
        include_tool_names=True,
        auto_save_context=True,
        max_tokens=None
    )
    
    # 创建代理
    agent = Agent(
        llm_service=llm,
        tool_registry=tools,
        user_resolver=user_resolver,
        agent_memory=agent_memory,
        llm_context_enhancer=ProductionLlmContextEnhancer(
            mysql_runner=db_tool.sql_runner, 
            agent_memory=agent_memory
        ),
        config=config
    )
    
    # 动态生成并缓存业务表建表语句到 business_contexts（参考 tests/test5.py 的文档入库方式）
    # 这一步是为了将表结构信息预加载到向量数据库中，提高后续查询的准确性
    try:
        enhancer: ProductionLlmContextEnhancer = agent.llm_context_enhancer  # type: ignore
        contexts = await enhancer.generate_business_contexts(COMMON_TABLES)
        # 保存到模块级变量，供其它测试/脚本使用
        global business_contexts
        business_contexts = contexts
        logger.info(f"业务建表语句已生成并缓存，共 {len(business_contexts)} 条")
        
        # 将建表语句以文本形式写入代理内存（与 tests/test5.py 风格一致）
        # 这样可以让LLM在生成SQL时参考这些结构化信息
        admin_user = User(id="admin", email="admin@example.com", group_memberships=['admin'])
        ctx = ToolContext(
            user=admin_user,
            conversation_id="business_schema_bootstrap",
            request_id="bootstrap_request",
            agent_memory=agent.agent_memory
        )
        for text in business_contexts:
            await agent.agent_memory.save_text_memory(content=text, context=ctx)
            # 展示缓存数据（截取前100字符避免日志过长）
            logger.info(f"已缓存建表语句: {text[:100]}...")
        logger.info("业务建表语句已写入代理内存")
    except Exception as e:
        logger.warning(f"生成或写入业务建表语句失败: {e}")
    
    logger.info("生产环境代理初始化完成")
    return agent


async def create_app():
    """
    创建应用实例，用于uvicorn启动
    
    Returns:
        FastAPI应用实例
    """
    # 创建代理
    agent = await create_production_agent()
    
    # 创建并返回FastAPI应用
    server = VannaFastAPIServer(agent=agent)
    app = server.create_app()
    logger.info("FastAPI应用实例已创建，可用于uvicorn启动")
    
    return app


def create_app_sync():
    """
    同步包装函数，用于uvicorn工厂模式
    
    解决uvicorn期望同步函数但我们需要异步初始化的问题
    使用nest_asyncio允许在已有事件循环中运行新的事件循环
    
    Returns:
        FastAPI应用实例
    """
    # 使用nest_asyncio运行异步函数
    import asyncio
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    
    app = loop.run_until_complete(create_app())
    return app


async def main():
    """主函数，启动生产环境服务器
    
    使用程序化方式启动uvicorn服务器，而不是命令行方式
    这样可以更好地控制服务器配置和错误处理
    """
    try:
        # 获取服务器配置
        server_config = ServerConfig.from_env()
        
        # 创建应用实例
        app = await create_app()
        
        # 创建并启动服务器
        logger.info(f"正在启动生产环境服务器，端口: {server_config.port}")
        
        # 获取uvicorn配置
        import uvicorn.config
        
        # 创建uvicorn配置
        config = uvicorn.config.Config(
            app=app,
            host="0.0.0.0",
            port=server_config.port,
            log_level="info"
        )
        
        # 创建服务器实例
        uvicorn_server = uvicorn.Server(config)
        
        # 启动服务器
        await uvicorn_server.serve()
        
    except Exception as e:
        logger.error(f"启动服务器时发生错误: {e}", exc_info=True)
        raise

if __name__ == "__main__":
    # 运行主函数
    # 使用asyncio.run启动异步主函数
    asyncio.run(main())
